{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HEiR6F5NE3V"
      },
      "source": [
        "# About this demo\n",
        "\n",
        "Colorization algorithms have a wide range of applications in various fields, such as historical photo restoration and artistic rendering.\n",
        "\n",
        "This demo shows the how to integrate the state of art AI algorithms into the BMF video processing pipeline. The famous open source colorization algorithm [DeOldify](https://github.com/jantic/DeOldify) is wrapped as an BMF pyhton module in less than 100 lines of codes.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oKr2leeRUDLN"
      },
      "source": [
        "# Acknowledgement\n",
        "The algorithm used in the demo is forked from the github project [DeOldify](https://github.com/jantic/DeOldify). We want to express our sincere appreciation and gratitude to the authors for this exceptional open-source project."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0kmtjxXwUQcA"
      },
      "source": [
        "# Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8oug66l9UadL"
      },
      "source": [
        "## 1. Environment Setup\n",
        "\n",
        "*   install the environment for the deoldify algorithm\n",
        "*   install the BMF\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "stS_u14PVE0i"
      },
      "source": [
        "### 1-1 git clone the deoldify project.\n",
        "The project is forked from the original github project. We add some suppport for Pillow Image input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yzqhgSOyygO4"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/eefengwei/DeOldify.git DeOldify"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o2-Ib9seVcwJ"
      },
      "source": [
        "### 1-2 install the dependent python packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5UKxle6h1daR"
      },
      "outputs": [],
      "source": [
        "!pip3 install -r ./DeOldify/requirements-colab.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t4SsR6VkVvGR"
      },
      "source": [
        "### 1-3 download the pretrained-weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O35abCsl5iw9"
      },
      "outputs": [],
      "source": [
        "!mkdir -p ./DeOldify/models\n",
        "!wget -c https://data.deepai.org/deoldify/ColorizeVideo_gen.pth -O ./DeOldify/models/ColorizeVideo_gen.pth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-IVElGWJWbqD"
      },
      "source": [
        "### 1-4 pip install BMF packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IjiiVnD9Aua-"
      },
      "outputs": [],
      "source": [
        "!pip install BabitMF-GPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KVQKIPamWnPA"
      },
      "source": [
        "### 1-5 verify the FFmpeg libraries is installed and version is correct\n",
        "\n",
        "BMF framework utilizes the FFmpeg video decoders and encoders as the built-in modules for video decoding and encoding. **It's neccessary for users to install supported FFmpeg libraries before using BMF.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WiPmTkhfaXft"
      },
      "source": [
        "**Optional step**\n",
        "\n",
        "Install ffmpeg and related libraries. For this demo, we don't have to do this step, because ffmpeg libraries are already installed in the Google Colab environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ARXN39-VaQlg"
      },
      "outputs": [],
      "source": [
        "!sudo apt install ffmpeg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PcLbXWWSYMfQ"
      },
      "source": [
        "List the ffmpeg libraries. It is expected that the related libraries such libavcodec, libavformat are installed. The output should be showned as below:\n",
        "![image.png]()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GLnRGvl5W00c"
      },
      "outputs": [],
      "source": [
        "!dpkg -l | grep -i ffmpeg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcLVtZ55YeWB"
      },
      "source": [
        "Display the FFmpeg version. **BMF supports the FFmpeg verions above 4.2.x.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "27HbW50PZLRG"
      },
      "outputs": [],
      "source": [
        "!ffmpeg -version"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AuU2xFfsb3VV"
      },
      "source": [
        "### 1-6 install wurlitzer\n",
        "\n",
        "This package is installed to show the BMF C++ logs in the colab console, otherwise only python logs are printed. This step is not necessary if you're not in a Colab or iPython notebook environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VoXqpHLh5aBM"
      },
      "outputs": [],
      "source": [
        "!pip install wurlitzer\n",
        "%load_ext wurlitzer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hC4ZKWhOd7g6"
      },
      "source": [
        "## 2. BMF processing pipeline programming\n",
        "*   implement the customized BMF python module, which is the DeOldify video colorization algorithm.\n",
        "*   set up the BMF processing pipeline (decoding->colorizaiton->encoding) and run it.\n",
        "*  display the results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Q-x00orgL-S"
      },
      "source": [
        "Add the DeOldify folder to the python search path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZM7B0XoG_FPG"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "sys.path.insert(0, '/content/DeOldify')\n",
        "print(sys.path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3mvqVqAgViw"
      },
      "source": [
        "### 2-1 Customized BMF module\n",
        "\n",
        "Implement the customized BMF python module, which is the DeOldify video colorization algorithm."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KRrbmFpL1im1"
      },
      "outputs": [],
      "source": [
        "%%file py_deoldify_module.py\n",
        "\n",
        "import bmf\n",
        "import numpy as np\n",
        "from bmf import ProcessResult, Packet, Timestamp, VideoFrame\n",
        "import PIL\n",
        "import bmf.hmp as mp\n",
        "from bmf.lib._bmf import sdk\n",
        "\n",
        "from deoldify import device\n",
        "from deoldify.device_id import DeviceId\n",
        "import torch\n",
        "from deoldify.visualize import *\n",
        "import warnings\n",
        "\n",
        "debug = False\n",
        "\n",
        "class py_deoldify_module(bmf.Module):\n",
        "    def __init__(self, node, option=None):\n",
        "        print(f'py_deoldify_module init ...')\n",
        "        self.node_ = node\n",
        "        self.option_ = option\n",
        "        print(option)\n",
        "        warnings.filterwarnings(\"ignore\", category=UserWarning, message=\".*?Your .*? set is empty.*?\")\n",
        "\n",
        "        #NOTE:  This must be the first call in order to work properly!\n",
        "        #choices:  CPU, GPU0...GPU7\n",
        "        device.set(device=DeviceId.GPU0)\n",
        "\n",
        "        if not torch.cuda.is_available():\n",
        "            print('warning: GPU is not available, the computation is going to be very slow...')\n",
        "\n",
        "        weight_path=Path('/content/DeOldify')\n",
        "        if option and 'model_path' in option.keys():\n",
        "            model_path = option['model_path']\n",
        "            if not model_path:\n",
        "                print(f'model_path={model_path}')\n",
        "                weight_path=Path(model_path)\n",
        "\n",
        "        self.colorizer = get_stable_video_colorizer(weight_path)\n",
        "        self.idx = 0\n",
        "\n",
        "        print(f'py_deoldify_module init successfully...')\n",
        "\n",
        "\n",
        "    def process(self, task):\n",
        "        # iterate through all input queues to the module\n",
        "        idx = self.idx\n",
        "\n",
        "        for (input_id, input_queue) in task.get_inputs().items():\n",
        "            # get output queue\n",
        "            output_queue = task.get_outputs()[input_id]\n",
        "\n",
        "            while not input_queue.empty():\n",
        "                # get the earliest packet from queue\n",
        "                packet = input_queue.get()\n",
        "\n",
        "                # handle EOF\n",
        "                if packet.timestamp == Timestamp.EOF:\n",
        "                    output_queue.put(Packet.generate_eof_packet())\n",
        "                    task.timestamp = Timestamp.DONE\n",
        "\n",
        "                # process packet if not empty\n",
        "                if packet.timestamp != Timestamp.UNSET and packet.is_(VideoFrame):\n",
        "\n",
        "                    vf = packet.get(VideoFrame)\n",
        "                    dst_md = sdk.MediaDesc().pixel_format(mp.kPF_RGB24)\n",
        "                    np_vf = sdk.bmf_convert(vf, sdk.MediaDesc(), dst_md).frame().plane(0).numpy()\n",
        "\n",
        "                    # numpy to PIL\n",
        "                    image = Image.fromarray(np_vf.astype('uint8'), 'RGB')\n",
        "\n",
        "                    colored_image = self.colorizer.colorize_single_frame_from_image(image)\n",
        "\n",
        "                    if not colored_image:\n",
        "                        print(f'Fail to process the input image with idx = {idx}')\n",
        "                        continue\n",
        "\n",
        "                    if debug:\n",
        "                        input_name = f'video/bmf_raw/frame_{idx}.png'\n",
        "                        print(f'input_name = {input_name}')\n",
        "                        image.save(input_name)\n",
        "\n",
        "                        output_name = f'video/bmf_out/frame_{idx}.png'\n",
        "                        print(f'output_name = {output_name}')\n",
        "                        colored_image.save(output_name)\n",
        "\n",
        "                    self.idx = idx + 1\n",
        "                    out_frame_np = np.array(colored_image)\n",
        "                    rgb = mp.PixelInfo(mp.kPF_RGB24)\n",
        "                    frame = mp.Frame(mp.from_numpy(out_frame_np), rgb)\n",
        "\n",
        "                    out_frame = VideoFrame(frame)\n",
        "                    out_frame.pts = vf.pts\n",
        "                    out_frame.time_base = vf.time_base\n",
        "\n",
        "                    pkt = Packet(out_frame)\n",
        "                    pkt.timestamp = out_frame.pts\n",
        "\n",
        "                    output_queue.put(pkt)\n",
        "\n",
        "\n",
        "        return ProcessResult.OK\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cCDhCLQjQzTz"
      },
      "source": [
        "### Note: Use Colab GPU steps\n",
        "\n",
        "If you want to use the free GPU provided by Google Colab, please follow such steps:\n",
        "\n",
        "*   Click on the \"Change\" tab in the menu bar.\n",
        "*   Select \"Notebook settings\" from the dropdown menu.\n",
        "*   In the \"Hardware accelerator\" section, make sure \"GPU\" is selected.\n",
        "*   Click on the \"Save\" button.\n",
        "\n",
        "After saving the changes, your Colab notebook will be configured to run on a GPU. You can verify this by running the following code snippet in a code cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YzAwyNvpPbYJ"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "    print('GPU is available!')\n",
        "else:\n",
        "    print('GPU is not available!')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ID7ScUgbPpQx"
      },
      "outputs": [],
      "source": [
        "!nvcc -V"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UmkF2cplgmDO"
      },
      "source": [
        "### 2-2 BMF pipeline setup and execution\n",
        "\n",
        "input_video_path is the input video path, and output_video_path is the output colored video path. model_weight_path is the path where you download the model weights in step 1-3."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wti6gSKDA3yL"
      },
      "outputs": [],
      "source": [
        "import bmf\n",
        "import py_deoldify_module\n",
        "\n",
        "input_video_path = '/content/DeOldify/test_videos/test_video.mp4'\n",
        "output_video_path = 'colored_video_2.mp4'\n",
        "model_weight_path = '/content/DeOldify/'\n",
        "\n",
        "graph = bmf.graph()\n",
        "\n",
        "video = graph.decode({\"input_path\": input_video_path})\n",
        "\n",
        "output_video = video['video'].module('py_deoldify_module', option={\"model_path\": model_weight_path})\n",
        "\n",
        "bmf.encode(\n",
        "    output_video[0],\n",
        "    None,\n",
        "    {\"output_path\": output_video_path}\n",
        "    ).run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jW8qsu-Og68d"
      },
      "source": [
        "### 2-3 Display the processed video"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sQuyM8EQ_3wa"
      },
      "outputs": [],
      "source": [
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "\n",
        "def show_video(video_path, video_width = 800):\n",
        "\n",
        "  video_file = open(video_path, \"r+b\").read()\n",
        "\n",
        "  video_url = f\"data:video/mp4;base64,{b64encode(video_file).decode()}\"\n",
        "\n",
        "  return HTML(f\"\"\"<video width={video_width} controls><source src=\"{video_url}\"></video>\"\"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EPVVdQKRdSF8"
      },
      "outputs": [],
      "source": [
        "# show input video\n",
        "show_video(input_video_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g09dbMaLdSF8"
      },
      "outputs": [],
      "source": [
        "# show output video\n",
        "show_video(output_video_path)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
